import torch
import torchvision
from PIL import Image
import os
from models.net import SimpleCNN
import config
import matplotlib.pyplot as plt

def predict_validation_set():
    # 设备设置
    device = config.DEVICE

    # 加载最佳模型（这里以最后一个epoch的模型为例）
    model_path = os.path.join(config.MODEL_SAVE_PATH, "mnist_cnn_epoch_10.pth")
    model = SimpleCNN().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()

    # 图片预处理（与训练时一致）
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.1307,), (0.3081,))
    ])

    # 验证集路径
    val_dir = r"F:\pycharm_program\MNIST_project\mnist_project\data\predict_set"

    # 结果文件路径 - 修改为当前目录下生成
    results_file = "infer_result.txt"  # 直接生成在当前工作目录
    # 或者指定绝对路径：
    # results_file = r"D:\Python\pythonProject4\mnist_write\prediction_results.txt"

    with open(results_file, "w", encoding="utf-8") as f:
        f.write("文件路径\t\t\t\t原始数值\t预测数值\n")
        f.write("=" * 80 + "\n")

    total = 0
    correct = 0

    # 遍历验证集目录结构
    for digit in range(10):
        digit_dir = os.path.join(val_dir, str(digit))
        if not os.path.exists(digit_dir):
            continue

        # 获取该数字目录下所有图片
        image_files = [f for f in os.listdir(digit_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

        for filename in image_files:
            image_path = os.path.join(digit_dir, filename)
            true_label = digit  # 从目录名获取真实标签

            try:
                # 加载和预处理图片
                image = Image.open(image_path).convert('L')  # 确保为灰度图
                image_tensor = transform(image).unsqueeze(0).to(device)

                # 预测
                with torch.no_grad():
                    output = model(image_tensor)
                pred_label = output.argmax(1).item()

                # 统计结果
                total += 1
                if pred_label == true_label:
                    correct += 1

                # 写入结果 - 使用完整绝对路径
                with open(results_file, "a", encoding="utf-8") as f:
                    f.write(f"{image_path}\t{true_label}\t{pred_label}\n")

            except Exception as e:
                print(f"处理图片 {image_path} 出错: {str(e)}")
                with open(results_file, "a", encoding="utf-8") as f:
                    f.write(f"{image_path}\tError: {str(e)}\n")

    # 计算并保存准确率
    if total > 0:
        accuracy = correct / total * 100
        with open(results_file, "a", encoding="utf-8") as f:
            f.write("\n" + "=" * 80 + "\n")
            f.write(f"测试总数: {total}\n")
            f.write(f"正确预测: {correct}\n")
            f.write(f"准确率: {accuracy:.2f}%\n")

        print(f"预测完成，结果已保存到 {os.path.abspath(results_file)}")
    else:
        print("未找到可处理的图片文件")
    total += 1
    if pred_label == true_label:
        correct += 1

    # 写入结果
    with open(results_file, "a", encoding="utf-8") as f:
        f.write(f"{image_path}\t{true_label}\t{pred_label}\n")

    # 可视化预测结果
    plt.imshow(image, cmap='gray')  # 显示图像
    plt.title(f'Predicted: {pred_label}')  # 标注预测值
    plt.axis('off')  # 不显示坐标轴
    plt.show()  # 展示图像


if __name__ == "__main__":
    predict_validation_set()